package com.shizhen.aidemo.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.shizhen.aidemo.Entity.ChatEntity;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.*;
import org.springframework.data.redis.core.RedisTemplate;

import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 使用Redis存储聊天历史的内存实现类
 */
@Slf4j
@Component
public class ChatRedisMemory implements ChatMemory {

    // Redis键名前缀
    private static final String KEY_PREFIX = "chat:history:";
    // Redis模板，用于操作Redis
    private final RedisTemplate<String, Object> redisTemplate;

    public ChatRedisMemory(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 向指定会话添加消息
     *
     * @param conversationId 会话ID
     * @param messages 消息列表
     */
    @Override
    public void add(String conversationId, List<Message> messages) {
        String key = KEY_PREFIX + conversationId;
        List<ChatEntity> listIn = new ArrayList<>();
        for (Message message : messages) {
            // 分割消息文本，提取实际内容
            String[] strs = message.getText().split("</think>");
            String text = strs.length == 2 ? strs[1] : strs[0];

            ChatEntity ent = new ChatEntity();
            ent.setChatId(conversationId);
            ent.setType(message.getMessageType().getValue());
            ent.setText(text);
            listIn.add(ent);
        }
        // 将消息实体列表存入Redis，并设置过期时间
        redisTemplate.opsForList().rightPushAll(key,listIn.toArray());
        redisTemplate.expire(key, 30, TimeUnit.MINUTES);
    }

    /**
     * 获取指定会话的最近N条消息
     *
     * @param conversationId 会话ID
     * @param lastN 最近的消息数量
     * @return 消息列表
     */
    @Override
    public List<Message> get(String conversationId, int lastN) {
        String key =KEY_PREFIX + conversationId;
        Long size = redisTemplate.opsForList().size(key);
        if (size == null || size ==0) {
            return Collections.emptyList();
        }

        int start = Math.max(0,(int) (size - lastN));
        List<Object> listTmp = redisTemplate.opsForList().range(key, start, -1);
        List<Message> listOut = new ArrayList<>();
        ObjectMapper objectMapper = new ObjectMapper();
        for (Object o : listTmp) {
            ChatEntity chat = objectMapper.convertValue(o, ChatEntity.class);
            // 根据消息类型创建相应的Message对象
            if (MessageType.USER.getValue().equals(chat.getType())) {
                listOut.add(new UserMessage(chat.getText()));
            } else if (MessageType.ASSISTANT.getValue().equals(chat.getType())) {
                listOut.add(new AssistantMessage(chat.getText()));
            } else if (MessageType.SYSTEM.getValue().equals(chat.getType())) {
                listOut.add(new SystemMessage(chat.getText()));
            }
        }
        return listOut;
    }

    /**
     * 清除指定会话的消息历史
     *
     * @param conversationId 会话ID
     */
    @Override
    public void clear(String conversationId) {
        redisTemplate.delete(KEY_PREFIX + conversationId);
    }
}
